package com.zzl.spark.sql.UDF_UDAF;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.*;
import scala.actors.threadpool.Arrays;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.LongStream;

public class UDAFJAVA {

    public static void main(String[] args) {
        System.out.println(Runtime.getRuntime().totalMemory()/1024/1024);
        SparkConf conf = new SparkConf().setMaster("local[9]").setAppName("DataFrameOps");
        JavaSparkContext sc = new JavaSparkContext(conf);
        System.setProperty("hadoop.home.dir", "E:\\hadoop");
        sc.setLogLevel("WARN");
        SparkSession sqlContext = SparkSession.builder().config(conf).getOrCreate();
        List<String> list = new ArrayList<>();
        list.add("yarn");
        list.add("yarn");
        list.add("marry");
        list.add("jack");
        list.add("to    m");
        list.add("tom");
        list.add("tom");
        list.add("tom");
        LongStream.range(0,10000000).forEach((i)->{
            list.add("yarn");
        });
        JavaRDD<String> parallelize = sc.parallelize(list);

        JavaRDD<Row> map = parallelize.map(RowFactory::create);

        ArrayList<StructField> structFields = new ArrayList<>();
        structFields.add(DataTypes.createStructField("name",DataTypes.StringType,true));

        StructType structType = DataTypes.createStructType(structFields);

        Dataset<Row> dataFrame = sqlContext.createDataFrame(map, structType);

        dataFrame.registerTempTable("nameTable");
        sqlContext.udf().register("stringCount", new UserDefinedAggregateFunction() {
            @Override
            public StructType inputSchema() {
                return DataTypes.createStructType(Arrays.asList(new StructField[]{DataTypes.createStructField("str", DataTypes.StringType, true)}));
            }

            @Override
            public DataType dataType() {
                return DataTypes.IntegerType;
            }

            @Override
            public boolean deterministic() {
                return true;
            }

            @Override
            public void update(MutableAggregationBuffer buffer, Row input) {
                buffer.update(0,buffer.getInt(0)+1);
            }

            @Override
            public StructType bufferSchema() {
                return DataTypes.createStructType(Arrays.asList(new StructField[]{DataTypes.createStructField("bf", DataTypes.IntegerType, true)}));
            }

            @Override
            public void merge(MutableAggregationBuffer buffer1, Row buffer2) {
                buffer1.update(0,buffer1.getInt(0)+buffer2.getInt(0));
            }

            @Override
            public void initialize(MutableAggregationBuffer buffer) {
                buffer.update(0,0);
            }

            @Override
            public Object evaluate(Row buffer) {
                return buffer.getInt(0);
            }
        });

        sqlContext.sql("select name,stringCount(name) from nameTable group by name").show();
        sc.stop();
        sqlContext.stop();
    }
}
